import cv2
import numpy as np
import matplotlib.pyplot as plt
import heapq
import math
import time
class Node:
    def __init__(self, x, y, g=0, h=0, parent=None):
        self.x = x
        self.y = y
        self.g = g
        self.h = h
        self.f = g + h
        self.parent = parent
    def __lt__(self, other):
        return self.f < other.f

class AStar:
    def __init__(self, grid):
        self.grid = grid
        self.rows, self.cols = grid.shape
        # 八方向移动及代价
        self.dirs = [
            (1,0,1),( -1,0,1),(0,1,1),(0,-1,1),
            (1,1,math.sqrt(2)),(-1,-1,math.sqrt(2)),
            (1,-1,math.sqrt(2)),(-1,1,math.sqrt(2))
        ]

    def in_bounds(self, x, y):
        return 0 <= x < self.cols and 0 <= y < self.rows

    def is_passable(self, x, y):
        return self.in_bounds(x, y) and self.grid[y, x] != 0

    def heuristic(self, x1, y1, x2, y2):
        # 使用对角线距离（更贴近真实路径）
        dx, dy = abs(x1 - x2), abs(y1 - y2)
        return dx + dy + (math.sqrt(2) - 2) * min(dx, dy)

    def find_path(self, start, end):
        open_list = []
        start_node = Node(*start, 0, self.heuristic(*start, *end))
        heapq.heappush(open_list, (start_node.f, start_node))
        came_from = {}
        g_score = {start: 0}

        while open_list:
            _, current = heapq.heappop(open_list)
            if (current.x, current.y) == end:
                # 回溯路径
                path = []
                while current:
                    path.append((current.x, current.y))
                    current = current.parent
                return path[::-1]

            for dx, dy, cost in self.dirs:
                nx, ny = current.x + dx, current.y + dy
                if not self.is_passable(nx, ny):
                    continue
                tentative_g = g_score[(current.x, current.y)] + cost
                if (nx, ny) not in g_score or tentative_g < g_score[(nx, ny)]:
                    g_score[(nx, ny)] = tentative_g
                    h = self.heuristic(nx, ny, *end)
                    neighbor = Node(nx, ny, tentative_g, h, current)
                    heapq.heappush(open_list, (neighbor.f, neighbor))

        return []


# ========================
# 主程序
# ========================
img = cv2.imread("map.png", cv2.IMREAD_GRAYSCALE)

# 定义可通行区域（去掉 0 和 192）
walkable_values = {255, 64, 32, 96, 120, 130, 140, 220}
grid = np.where(np.isin(img, list(walkable_values)), 1, 0)

astar = AStar(grid)

start, end = None, None

def onclick(event):
    global start, end
    if event.xdata is None or event.ydata is None:
        return
    x, y = int(event.xdata), int(event.ydata)
    if not astar.is_passable(x, y):
        print("❌ 点击在障碍物上:", (x, y))
        return
    if start is None:
        start = (x, y)
        plt.scatter([x], [y], color="green", s=50)
        plt.draw()
        print("✅ 起点:", start)
    elif end is None:
        end = (x, y)
        plt.scatter([x], [y], color="blue", s=50)
        print("✅ 终点:", end)
        # ====== 计算路径并测时间 ======
        t0 = time.time()
        path = astar.find_path(start, end)
        t1 = time.time()
        elapsed = t1 - t0
        print(f"⏱ 路径计算耗时: {elapsed:.6f} 秒")
        if path:
            px, py = zip(*path)
            plt.plot(px, py, color="red")
            print("✅ 找到路径, 长度:", len(path))
        else:
            print("❌ 没有找到路径")
        plt.draw()

plt.imshow(grid, cmap="gray")
plt.title("点击选择起点 (绿色) 和终点 (蓝色)")
cid = plt.gcf().canvas.mpl_connect('button_press_event', onclick)
plt.show()
